package com.fyh.config;

import com.fyh.constant.JwtConstant;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.Jwts;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;

import java.security.Principal;
import java.util.Map;

@Component
@Slf4j
public class CustomHandshakeHandler extends DefaultHandshakeHandler {

    @Override
    protected Principal determineUser(ServerHttpRequest request,
                                      WebSocketHandler wsHandler,
                                      Map<String, Object> attributes) {
        log.info("开始处理WebSocket握手请求");

        // 从URL查询参数中获取token
        String jwtToken = extractTokenFromQueryString(request.getURI().getQuery());

        if (jwtToken != null && !jwtToken.isEmpty()) {
            // 解析JWT令牌获取userId
            Long userId = validateToken(jwtToken);
            if (userId == null) {
                log.warn("JWT验证失败");
                return null;
            }
            log.info("userId:{}", userId);
            attributes.put("userId", userId);
            return () -> userId.toString();
        } else {
            log.warn("未提供JWT令牌");
            return null;
        }
    }

    private String extractTokenFromQueryString(String queryString) {
        if (queryString == null || queryString.isEmpty()) {
            return null;
        }

        String[] params = queryString.split("&");
        for (String param : params) {
            if (param.startsWith("token=")) {
                return param.substring(6);
            }
        }
        return null;
    }

    private Long validateToken(String token) {
        try {
            log.info("开始验证JWT令牌：{}", token);
            Claims claims = Jwts.parser()
                    .setSigningKey(JwtConstant.SIGNATRUE)
                    .parseClaimsJws(token)
                    .getBody();
            return ((Integer) claims.get("userId")).longValue();
        } catch (Exception e) {
            log.error("JWT验证异常：", e);
            return null;
        }
    }
}
